import torch
import torch.nn.functional as F
import torch.optim as optim

from defense.base_defense import BaseDefense

class Defense(BaseDefense):
    def __init__(self, model_name, config):
        super().__init__('IGR', model_name, config, retrain = True)

    def train(self, loader, opt, epoch):
        self.model.train()
        for ind, (images, labels) in enumerate(loader):
            images.requires_grad = True
            images = images.to(self.model.device)
            labels = labels.to(self.model.device)

            output = self.model(images)
            loss_1 = F.cross_entropy(output, labels)

            grad = torch.autograd.grad(loss_1, images, create_graph=True)[0]
            loss_2 = torch.Tensor.norm(grad, p=2) ** 2
            total_loss =loss_1 + loss_2 * self.config['lambda_r']

            opt.zero_grad()
            total_loss.backward()
            opt.step()

            ind += 1
            print('\rEpoch {}: [batch: {}/{} ({:.0f}%)] \tloss_1:{:.4f} + lambda:316 * loss_2:{:.4f} = {:.4f}'.format(
                epoch, ind, len(loader), ind / len(loader) * 100, loss_1, loss_2, total_loss), end=' ')

        print()

    def defend(self):
        batch_size = self.model.training_params['batch_size']
        # 加载训练数据
        train_loader, valid_loader = self.model.get_train_valid_loader(batch_size=batch_size,
                                                                       valid_size=0.1,
                                                                       transform_train=True,
                                                                       shuffle=True)
        # 使用Adam优化函数
        opt = optim.Adam(self.model.parameters(), lr=0.0002, eps=1e-4)

        best_ratio = 0
        best_state_dict = None
        for epoch in range(self.model.training_params['num_epochs']):
            self.train(train_loader, opt, epoch)
            ratio = self.model.test(valid_loader)
            if ratio > best_ratio:
                best_ratio = ratio
                best_state_dict = self.model.state_dict_clone()  # 保存此时网络参数
                print('Validate: Accuracy increased')
            else:
                print('Validate: Accuracy decreased')
            print()
        return best_state_dict